package com.zcyr.core.sql.visitor;

import com.alibaba.druid.sql.ast.SQLExpr;
import com.alibaba.druid.sql.ast.SQLObject;
import com.alibaba.druid.sql.ast.expr.*;
import com.alibaba.druid.sql.ast.statement.*;
import com.alibaba.druid.sql.parser.Token;
import com.zcyr.core.sql.statement.OGCreateSchemaStatement;
import com.zcyr.core.sql.statement.OGDropSchemaStatement;

import java.util.HashMap;
import java.util.Map;

/**
 * 实现对sql的特殊处理，起到隔离用户操作的作用
 */
public class OGTrimVisitor extends OGASTVisitorAdapter{
    private final String suffix;
    // opengauss保留字
    private final Map<String, Token> keywords = new HashMap<>();

    public OGTrimVisitor(String suffix) {
        super();
        this.suffix = suffix;
        keywords.put("FALSE", Token.FALSE);
        keywords.put("INTO", Token.INTO);
    }

    /**
     * 判断表达式，如果不是是opengauss保留字则返回true
     */
    private boolean isNotKeyword(SQLObject expr){
        if (expr instanceof SQLIdentifierExpr
                && keywords.get(((SQLIdentifierExpr) expr).getName().toUpperCase()) != null){
            return false;
        }
        if (expr instanceof SQLAllExpr){
            return false;
        }
        return !(expr instanceof SQLQueryExpr);
    }

    private String wrapString(String target){
        if (target.isBlank()){
            return target;
        }
        return target + suffix;
    }

    @Override
    public boolean visit(OGCreateSchemaStatement x) {
        if (isNotKeyword(x.getSchema())){
            SQLIdentifierExpr t = (SQLIdentifierExpr) x.getName();
            t.setName(wrapString(t.getName()));
        }
        return super.visit(x);
    }

    @Override
    public boolean visit(OGDropSchemaStatement x) {
        return false;
    }

    @Override
    public boolean visit(SQLExprTableSource x){
        if (x.getSchema() != null){
            x.setSchema(wrapString(x.getSchema()));
        }
        return super.visit(x);
    }

    @Override
    public boolean visit(SQLSelectItem x){
        SQLExpr expr = x.getExpr();
        if (expr instanceof SQLPropertyExpr){
            expr = ((SQLPropertyExpr) expr).getOwner();
            if (expr instanceof SQLPropertyExpr){
                ((SQLPropertyExpr) expr).setOwner(wrapString(((SQLPropertyExpr) expr).getOwnernName()));
            }
        }
        return super.visit(x);
    }

    @Override
    public boolean visit(SQLBinaryOpExpr x){
        SQLExpr expr = x.getLeft();
        if (expr instanceof SQLPropertyExpr){
            expr = ((SQLPropertyExpr) expr).getOwner();
            if (expr instanceof SQLPropertyExpr){
                ((SQLPropertyExpr) expr).setOwner(wrapString(((SQLPropertyExpr) expr).getOwnernName()));
            }
        }
        expr = x.getRight();
        if (expr instanceof SQLPropertyExpr){
            expr = ((SQLPropertyExpr) expr).getOwner();
            if (expr instanceof SQLPropertyExpr){
                ((SQLPropertyExpr) expr).setOwner(wrapString(((SQLPropertyExpr) expr).getOwnernName()));
            }
        }

        return super.visit(x);
    }

}
